#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <cmath>
#include "Matrix.h"
#include "KalmanFilter.h"

using namespace std;

// Function for Kalman filter prediction
void cpp(double x, double y, double t, KalmanFilter& kf, double& t_last, vector<double>& x_pred, vector<double>& y_pred, vector<double>& t_pred) {
    kf.updateStateTransitionMatrix(t - t_last);
    kf.predict();

    Vector2 z;
    z(0) = x;
    z(1) = y;
    kf.update(z);

    for (int i = 0; i < 25; ++i) {
        kf.updateStateTransitionMatrix(i);
        kf.extrapolate();
        Vector4 predicted_pos = kf.getStateEx();
        x_pred.push_back(predicted_pos(0));
        y_pred.push_back(predicted_pos(1));
        t_pred.push_back(t + i);
    }
    t_last = t;
}

int main() {
    string file = "input.txt";
    string savepath = "D:/";

    ifstream input_file(file);
    vector<string> data_lines;
    string line;

    while (getline(input_file, line)) {
        data_lines.push_back(line);
    }

    int len = data_lines.size();
    vector<double> time(len), abs_x(len), abs_y(len);

    for (int i = 0; i < len; ++i) {
        size_t first_col = data_lines[i].find(":");
        size_t second_col = data_lines[i].find(":", first_col + 1);
        time[i] = stod(data_lines[i].substr(first_col + 1, second_col - first_col - 1)) * 60 +
            stod(data_lines[i].substr(second_col + 1)) * 1000;
        time[i] -= time[0];  // convert to ms

        abs_x[i] = stod(data_lines[i].substr(9, data_lines[i].find("\t", 9) - 9));
        abs_x[i] -= abs_x[0];
        abs_y[i] = stod(data_lines[i].substr(12, data_lines[i].find("\t", 12) - 12));
        abs_y[i] -= abs_y[0];
    }

    // Define refresh frequencies to iterate over
    vector<int> frequencies = { 120, 60, 30 };

    for (int f : frequencies) {
        cout << f << " Hz" << endl;

        // Generate refreshed cursor position (baseline)
        int refresh_time_inc = round(1000 / f);
        int refresh_len = static_cast<int>(time.back() / refresh_time_inc) + 1;
        vector<double> refresh_time(refresh_len);
        vector<int> refresh_index(refresh_len);

        for (int i = 0; i < refresh_len; ++i) {
            refresh_time[i] = i * refresh_time_inc;
            refresh_index[i] = 0;
        }

        for (int i = 1; i < refresh_len; ++i) {
            for (int k = refresh_index[i - 1]; k < len; ++k) {
                if (time[k] > refresh_time[i]) {
                    refresh_index[i] = k - 1;
                    break;
                }
            }
        }

        vector<double> refresh_abs_x(refresh_len), refresh_abs_y(refresh_len);
        for (int i = 0; i < refresh_len; ++i) {
            refresh_abs_x[i] = abs_x[refresh_index[i]];
            refresh_abs_y[i] = abs_y[refresh_index[i]];
        }

        // Kalman filter initialization
        KalmanFilter kf(10, 0.01);
        double t_last = 0;
        vector<double> refresh_abs_x_ex = refresh_abs_x;
        vector<double> refresh_abs_y_ex = refresh_abs_y;

        for (int n = 0; n < refresh_index.back() + 1; ++n) {
            vector<double> x_pred, y_pred, t_pred;
            cpp(abs_x[n], abs_y[n], time[n], kf, t_last, x_pred, y_pred, t_pred);

            int last_index = static_cast<int>(time[n] / refresh_time_inc);
            for (int m = 0; m < 25; ++m) {
                if (t_pred[m] >= time[n + 1]) {
                    break;
                }
                if (t_pred[m] >= refresh_time[last_index + 1] - 1) {
                    refresh_abs_x_ex[last_index + 1] = x_pred[m];
                    refresh_abs_y_ex[last_index + 1] = y_pred[m];
                    break;
                }
            }
        }

        // Save data
        ofstream out_file;
        out_file.open(savepath + "refresh_abs_x_" + to_string(f) + "_Hz_perfect_circle.txt");
        for (double val : refresh_abs_x) {
            out_file << val << endl;
        }
        out_file.close();

        out_file.open(savepath + "refresh_abs_y_" + to_string(f) + "_Hz_perfect_circle.txt");
        for (double val : refresh_abs_y) {
            out_file << val << endl;
        }
        out_file.close();

        out_file.open(savepath + "refresh_abs_x_ex_" + to_string(f) + "_Hz_perfect_circle.txt");
        for (double val : refresh_abs_x_ex) {
            out_file << val << endl;
        }
        out_file.close();

        out_file.open(savepath + "refresh_abs_y_ex_" + to_string(f) + "_Hz_perfect_circle.txt");
        for (double val : refresh_abs_y_ex) {
            out_file << val << endl;
        }
        out_file.close();

        out_file.open(savepath + "refresh_time_" + to_string(f) + "_Hz_perfect_circle.txt");
        for (double val : refresh_time) {
            out_file << val << endl;
        }
        out_file.close();
    }

    return 0;
}